﻿using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;

namespace DotNetDetour.DetourWays
{
    public class Memory : IDisposable
    {
        private List<MemoryPool> _pools = new List<MemoryPool>();

        private IntPtr calc_align_up_bit(IntPtr p, long align)
        {
            return (IntPtr)(((long)p + align - 1) & (~(align - 1)));
        }

        private IntPtr calc_align_down_bit(IntPtr p, long align)
        {
            return (IntPtr)((long)p & (~(align - 1)));
        }

        private MemoryPool FindPool(IntPtr nearby)
        {
            NativeAPI.GetSystemInfo(out SYSTEM_INFO si);
            ////0x00007ffffffeffff
            //var maxAddr = si.lpMaximumApplicationAddress;
            ////0x0000000000010000
            //var minAddr = si.lpMinimumApplicationAddress;

            var maxAddr = (IntPtr)Math.Min(
              (long)calc_align_up_bit(nearby + 5 + int.MaxValue, si.dwAllocationGranularity),
              (long)si.lpMaximumApplicationAddress);
            var minAddr = (IntPtr)Math.Max(
              (long)calc_align_down_bit(nearby + 5 + int.MinValue, si.dwAllocationGranularity),
              (long)si.lpMinimumApplicationAddress);

            foreach (var item in _pools)
            {
                if ((long)item.Start > (long)minAddr && (ulong)item.End < (ulong)maxAddr)
                    return item;
            }

            IntPtr buffPtr = IntPtr.Zero;
            if (IntPtr.Size == 8)
            {
                foreach (int i in new int[] { -1, 1 })
                {
                    var bAdrr = nearby;
                    while (0 != NativeAPI.VirtualQuery(bAdrr, out MEMORY_BASIC_INFORMATION64 mbi,
                        Marshal.SizeOf(typeof(MEMORY_BASIC_INFORMATION64))))
                    {
                        if (mbi.State == NativeAPI.MEM_FREE)
                        {
                            buffPtr = NativeAPI.VirtualAlloc(mbi.BaseAddress, 0x1000,
                              NativeAPI.MEM_RESERVE | NativeAPI.MEM_COMMIT,
                              Protection.PAGE_EXECUTE_READWRITE);

                            goto LnBreak;
                        }
                        bAdrr = (IntPtr)((long)mbi.AllocationBase + si.dwAllocationGranularity * i);
                        if ((long)bAdrr > (long)maxAddr || (long)bAdrr < (long)minAddr) break;
                    }
                }
            LnBreak:;
            }
            else
            {
                buffPtr = NativeAPI.VirtualAlloc(IntPtr.Zero, si.dwAllocationGranularity,
                  NativeAPI.MEM_COMMIT | NativeAPI.MEM_RESERVE,
                  Protection.PAGE_EXECUTE_READWRITE);
            }

            if (buffPtr == IntPtr.Zero)
                throw new InsufficientMemoryException();

            var pool = new MemoryPool(buffPtr, si.dwAllocationGranularity, 32);
            _pools.Add(pool);
            return pool;
        }

        public IntPtr Alloc(IntPtr nearby)
        {
            var pool = FindPool(nearby);
            if (pool == null)
                return IntPtr.Zero;

            return pool.Alloc();
        }

        public void Free(IntPtr p)
        {
            foreach (var item in _pools)
            {
                if (item.Free(p))
                    break;
            }
        }

        public void Dispose()
        {
            foreach (var item in _pools)
            {
                NativeAPI.VirtualFree(item.Start, 0, NativeAPI.MEM_RELEASE);
            }
            _pools.Clear();
        }
    }

    public class MemoryPool
    {
        public IntPtr Start;
        public IntPtr End;

        public uint Size
        {
            get
            {
                return (uint)((long)End - (long)Start);
            }
        }

        protected int _blocksize;
        private int _freeIdx;

        public MemoryPool(IntPtr buffer, uint bufferlen, int blocksize)
        {
            Initial(buffer, bufferlen, blocksize);
        }

        public void Initial(IntPtr buffer, uint bufferlen, int blocksize)
        {
            Start = buffer;
            End = buffer + (int)bufferlen;
            _blocksize = blocksize;

            Marshal.WriteInt32(Start, -1);
            _freeIdx = 0;
        }

        public IntPtr Alloc()
        {
            var freePtr = Start + _freeIdx * (_blocksize + 4);

            if ((long)freePtr + (_blocksize + 4) > (long)End)
                throw new InsufficientMemoryException();

            IntPtr ret;

            var idx = Marshal.ReadInt32(freePtr);
            if (idx == -1)
            {
                ret = freePtr + 4;
                freePtr += (_blocksize + 4);
                Marshal.WriteInt32(freePtr, -1);
                _freeIdx++;
            }
            else
            {
                _freeIdx = idx;
                ret = freePtr + 4;
            }
            return ret;
        }

        public bool Free(IntPtr p)
        {
            if ((ulong)p > (ulong)End || (ulong)p < (ulong)Start)
                return false;

            var freePtr = p - 4;
            Marshal.WriteInt32(freePtr, _freeIdx);

            _freeIdx = (int)((ulong)freePtr - (ulong)Start) / (_blocksize + 4);

            return true;
        }
    }
}